import os
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
from pyspark.sql import SparkSession

# 创建 SparkSession
spark = SparkSession.builder.appName('NBAPlayerStatsAnalysis2').getOrCreate()

# 数据路径配置
data_path = "/home/developer/NBA_Bigdata/数据/cleaned/merged_cleaned_data.csv"
output_image_path = "/home/developer/image/pearson.png"

# 确保图片保存路径存在
os.makedirs(os.path.dirname(output_image_path), exist_ok=True)

# 加载数据
try:
    print("加载数据中...")
    data = spark.read.option("header", True).option("inferSchema", True).csv(data_path)
    print("数据加载成功！")
except Exception as e:
    print(f"数据加载失败：{e}")
    exit(1)

# 将 Spark DataFrame 转换为 Pandas DataFrame
print("转换数据为 Pandas DataFrame...")
try:
    df = data.toPandas()
    print("转换成功！")
except Exception as e:
    print(f"数据转换失败：{e}")
    exit(1)

# 检查列的数据类型
print("检查列数据类型：")
print(df.dtypes)

# 筛选出数值列
print("筛选数值列...")
numeric_df = df.select_dtypes(include=['number'])

# 检查筛选结果
if numeric_df.empty:
    print("未找到数值列，无法计算相关性矩阵！")
    exit(1)
else:
    print("筛选后的数值数据：")
    print(numeric_df.head())

# 计算相关性矩阵
print("计算相关性矩阵...")
corr = numeric_df.corr()

# 生成热图
print("生成相关性矩阵热图...")
try:
    # 生成遮罩，遮住上三角矩阵
    mask = np.triu(np.ones_like(corr, dtype=bool))

    # 设置 matplotlib 图像尺寸
    f, ax = plt.subplots(figsize=(11, 9))

    # 定义自定义的调色板
    cmap = sns.diverging_palette(230, 20, as_cmap=True)

    # 绘制热图
    sns.heatmap(corr, mask=mask, cmap=cmap, vmax=.3, center=0,
                square=True, linewidths=.5, cbar_kws={"shrink": .5})

    # 保存图片
    plt.savefig(output_image_path)
    print(f"相关性矩阵热图已保存到 {output_image_path}")
except Exception as e:
    print(f"生成或保存热图失败：{e}")

# 显示图像（可选）
# plt.show()

print("分析完成！")
